import os

import numpy as np
import pandas as pd
from sqlalchemy import create_engine
# from sklearn.preprocessing import MinMaxScaler as tran
from sklearn.metrics import mean_absolute_error as mae
from tqdm import tqdm
from matplotlib.figure import Figure
from matplotlib.ticker import MultipleLocator
from matplotlib.backends.backend_pdf import PdfPages
from pandas import DataFrame
import tushare as ts

ts.set_token("3cea78a214f1ed12982e81dfc42750e149f12c7c7d99e180633982d4")
pro = ts.pro_api()

db_future_main = 'mysql+pymysql://hf_user:290202@192.168.3.179:3306/future?charset=utf8&use_unicode=1'
engine_future_main = create_engine(db_future_main)

db_future = 'mysql+pymysql://hf_user:290202@192.168.3.179:3306/future_short?charset=utf8&use_unicode=1'
engine_future = create_engine(db_future)

db_stock = 'mysql+pymysql://hf_user:290202@192.168.3.179:3306/stock?charset=utf8&use_unicode=1'
engine_stock = create_engine(db_stock)


class Tran:
    def __init__(self):
        super(Tran, self).__init__()

    @staticmethod
    def fit_transform(df: DataFrame):
        assert 'close' in df.columns.to_list()
        v = df['close'].to_list()[0]
        return df / v


# def mae(xx, yy):
#     xx_start = xx.iloc[0, 0]
#     xx_end = xx.iloc[-1, 0]
#     yy_start = yy.iloc[0, 0]
#     yy_end = yy.iloc[-1, 0]
#     return abs(yy_end - xx_end) + abs(yy_start - xx_start)


def find_tables(db_):
    tbs = pd.read_sql("show tables", db_).iloc[:, 0].to_list()
    print(f"total {len(tbs)} tables")
    return tbs


def load_df_by_tick_and_db(tick, db):
    map_db = {'future': engine_future, 'stock': engine_stock}
    df = pd.read_sql(f"select * from {tick}", map_db[db])
    return df


def load_target_data(code="SF.ZCE"):
    path = f"cache/{code}.csv"
    if os.path.exists(path):
        data = pd.read_csv(path, index_col=0)
    else:
        data = pro.fut_daily(ts_code=code, )
        data = data.sort_values(by=['trade_date'], ascending=True)
        data.to_csv(path)
    return data


def find_similar(target_df, match_nums):
    target_length = len(target_df)
    metrics = []

    mm_for_y = Tran()
    mm_y = mm_for_y.fit_transform(target_df)

    tbs_for_match = find_tables(db_future)
    n_for_match = min(match_nums, len(tbs_for_match))
    pbar = tqdm(tbs_for_match[:n_for_match], desc='Tick', ncols=80)

    for tick in pbar:
        try:
            x_data = load_df_by_tick_and_db(tick, 'future')
            trade_dates = x_data.trade_date
            x_data = x_data.loc[:, target_df.columns]

            for i in range(len(x_data) - target_length - 10):
                temp_x_data = x_data.iloc[i + 4:i + 4 + target_length]
                t = trade_dates[i + 4]
                mm_for_x = Tran()
                mm_x = mm_for_x.fit_transform(temp_x_data)
                xy_mae = mae(mm_y, mm_x)
                se = pd.Series({'tick': tick, 'database': 'future', 't': t, 'mae': xy_mae})
                metrics.append(se)
            pbar.set_postfix(tick=tick)
        except:
            pass
    return pd.DataFrame(metrics)


def plot_fig(y_se, r_se, i):
    tick = r_se['tick']
    t = r_se['t']
    db = r_se['database']

    fig = Figure(figsize=(20, 9), dpi=100)
    ax = fig.add_subplot(111)

    data = load_df_by_tick_and_db(tick, db)
    data_fit = data[data.trade_date > t]
    data_fit.index = list(range(len(data_fit)))

    fit_close = data_fit.close
    reset_ratio = fit_close.to_list()[0] / y_se.to_list()[0]
    fit_close = fit_close / reset_ratio

    t_end = len(y_se) - 1
    pred_se = pd.Series(index=list(range(len(y_se) - 1, len(data_fit))), dtype='float')
    pred_se[pred_se.index.to_list()[0]] = y_se.to_list()[-1]
    for pred_idx in pred_se.index.to_list()[1:]:
        pred_se[pred_idx] = pred_se[pred_idx - 1] + (fit_close[pred_idx] - fit_close[pred_idx - 1])

    ax.plot(fit_close, label=f'history similar {tick}')
    ax.plot(y_se, label='y')
    ax.plot(pred_se, label='pred y with similar')
    ax.vlines(x=t_end, color='red', ymin=min(fit_close), ymax=max(fit_close))
    ax.xaxis.set_major_locator(MultipleLocator(12))
    ax.set_title(f'i_{i}_tick_{tick}_t{t}.png')
    ax.legend()
    return fig, pred_se, fit_close


def remove_number_get_tk(ss):
    for i in range(10):
        ss = ss.replace(str(i), '')
    return ss.split('_')[0]


def gen_pdf(number, v_df, y_se, code):
    ticks, se_li, tks = [], [], []
    assert number < len(v_df)
    for i, r_se in v_df.iloc[:number, :].iterrows():
        tk = remove_number_get_tk(r_se['tick'])
        if tk not in tks:
            tks.append(tk)
            if r_se['tick'] not in ticks:
                fg, pred_se, fit_close = plot_fig(y_se, r_se, i)
                se_li.append(pred_se)
                ticks.append(r_se['tick'])
    _, _, path = generate_band(fit_close, se_li, code)
    return path


def generate_band(fit_close, se_li, code):
    fig = Figure(figsize=(20, 9), dpi=100)
    ax = fig.add_subplot(111)
    ax.plot(fit_close, label='y--', color='white')
    for se in se_li:
        ax.plot(se)
    v_df = pd.concat(se_li, axis=1)
    boll_df = pd.DataFrame(columns=['v_25', 'v_50', 'v_75'])
    if len(v_df) > 200:
        v_df = v_df.iloc[:199, :]
    for r, r_se in v_df.iterrows():
        r_se = r_se.dropna()
        if len(r_se) > 5:
            r_se.sort_values(ascending=True, inplace=True)
            v_75 = np.quantile(r_se, 0.75)
            v_50 = np.quantile(r_se, 0.5)
            v_25 = np.quantile(r_se, 0.25)
            append_df = pd.DataFrame([v_25, v_50, v_75], columns=[r], index=['v_25', 'v_50', 'v_75'])
            boll_df = boll_df.append(append_df.T)
    fig_v = Figure(figsize=(20, 9), dpi=100)
    ax_v = fig_v.add_subplot(111)
    for name in ['v_25', 'v_50', 'v_75']:
        ax_v.plot(boll_df[name], label=f'quantile_{name}')
    ax_v.legend()
    boll_df.to_excel(f"cache/{code}_boll_df.xlsx")
    return fig, fig_v, f"cache/{code}_boll_df.xlsx"


def get_listed_tick():
    """
    获取future主力合约库中所有的表名
    :return: list
    """
    tks = pd.read_sql("show tables", engine_future_main)
    return tks.iloc[:, 0].to_list()


def auto_pred(code):
    similar_days = 60  # 回测天数（从现在开始）
    similar_cols = ['close', 'open', 'high', 'low']  # 计算相似性的列名
    match_nums = 1000  # 去期货数据库匹配指标的数量，大于表的数量
    top_nums = 2000  # 匹配后取top数量

    dt = load_target_data(code).loc[:, similar_cols]  # 获取待预测的指标
    dt = dt.iloc[-similar_days:, :]  # 根据天数切分数据
    sorting_mae_df = find_similar(dt, match_nums).sort_values(by='mae', ascending=True)  # 匹配生产特征表
    dt.index = list(range(len(dt)))
    y_se = dt.close
    path = gen_pdf(top_nums, sorting_mae_df, y_se, code)  # 生成可能的轨迹，并对轨迹求百分位
    return path


if __name__ == "__main__":
    code = "JM.DCE"
    auto_pred(code)

